{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Patents Document Understanding with Gemini\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fdocument-processing%2Fpatents_understanding.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/document-processing/patents_understanding.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/document-processing/patents_understanding.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) |  [Holt Skinner](https://github.com/holtskinner) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "Back in 2019, I wrote a [Google Cloud Blog post](https://cloud.google.com/blog/products/ai-machine-learning/building-a-document-understanding-pipeline-with-google-cloud) in collaboration with [Michael Munn](https://github.com/munnm) and [Michael Sherman](https://github.com/michaelwsherman) which illustrates how to build a Document Understanding Pipeline using [AutoML](https://cloud.google.com/automl).\n",
        "\n",
        "This example showed how to train custom machine learning models for the following tasks:\n",
        "\n",
        "- [Image Classification](https://cloud.google.com/vision/automl/docs/beginners-guide)\n",
        "- [Entity Extraction](https://cloud.google.com/natural-language/automl/entity-analysis/docs/)\n",
        "- [Text Classification](https://cloud.google.com/natural-language/automl/docs/predict)\n",
        "- [Object Detection](https://cloud.google.com/vision/automl/object-detection/docs/)\n",
        "\n",
        "In today's world of Generative AI models like [Gemini](https://blog.google/technology/ai/google-gemini-ai/), it's possible to create the same document processing pipeline without training custom models. This significantly simplifies the process and reduces the time & resources required to automate these workflows.\n",
        "\n",
        "In this notebook, we'll create a document understanding pipeline on a public dataset of [patents PDFs](https://console.cloud.google.com/marketplace/details/global-patents/labeled-patents) stored in BigQuery and use [Batch Prediction for Gemini 2.0 in Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini) to:\n",
        "\n",
        "- Classify the patent granter (US or EU).\n",
        "- Classify the invention type (Medical Tech, Computer Vision, Cryptography, Other).\n",
        "- Extract key entities like publication date, application number, etc.\n",
        "- Detect bounding boxes for figures in the document."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Google Gen AI SDK for Python\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-genai google-cloud-bigquery pandas-gbq"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. In Colab or Colab Enterprise, you might see an error message that says \"Your session crashed for an unknown reason.\" This is expected. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8aMIcn9mEWt"
      },
      "source": [
        "### Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rrgrbhPmmEWt"
      },
      "outputs": [],
      "source": [
        "from google import genai\n",
        "from google.cloud import bigquery\n",
        "import pandas_gbq"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and create client\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)\n",
        "bq_client = bigquery.Client()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e43229f3ad4f"
      },
      "source": [
        "### Load the Gemini 2.0 Flash model\n",
        "\n",
        "To learn more about all [Gemini models on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-models)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf93d5f0ce00"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.0-flash-001\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EdvJRUWRNGHE"
      },
      "source": [
        "## The Prediction Pipeline\n",
        "\n",
        "Now, let's build the pipeline for processing the patent documents. We will:\n",
        "\n",
        "1. Fetch the PDF URIs from BigQuery into a Pandas DataFrame.\n",
        "2. Create a detailed prompt for Gemini to extract the information.\n",
        "3. Process each PDF with the prompt and save the results in a Pandas DataFrame\n",
        "4. Save the structured output to a new BigQuery table"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aa191bc13022"
      },
      "source": [
        "### Get PDF URIs from BigQuery\n",
        "\n",
        "We'll query BigQuery to get a list of Google Cloud Storage URIs for the patent PDF files.\n",
        "\n",
        "NOTE: This query limits to only 5 documents to save on processing time for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mzVYcJ5CYgZ3"
      },
      "outputs": [],
      "source": [
        "query = \"\"\"\n",
        "SELECT\n",
        "  COALESCE(t1.gcs_path, t2.gcs_path, t3.gcs_path) AS gcs_path\n",
        "FROM\n",
        "  `bigquery-public-data.labeled_patents.extracted_data` AS t1\n",
        "FULL OUTER JOIN\n",
        "  `bigquery-public-data.labeled_patents.figures` AS t2\n",
        "ON\n",
        "  t1.gcs_path = t2.gcs_path\n",
        "FULL OUTER JOIN\n",
        "  `bigquery-public-data.labeled_patents.invention_types` AS t3\n",
        "ON\n",
        "  COALESCE(t1.gcs_path, t2.gcs_path) = t3.gcs_path\n",
        "LIMIT 5\n",
        "\"\"\"\n",
        "\n",
        "df = bq_client.query(query).result().to_dataframe()\n",
        "df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ea36f19d9f7"
      },
      "source": [
        "### Define the Gemini prompt\n",
        "\n",
        "Here's the prompt we'll use with Gemini to extract the information we need. It's a detailed instruction that specifies the output format as JSON and also includes a JSON Schema."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jYxX_rCGe0hC"
      },
      "outputs": [],
      "source": [
        "PATENTS_PROMPT = \"\"\"\n",
        "Given a patent document, please extract the following information and output it as a JSON object with the specified keys.\n",
        "\n",
        "**JSON Output Format:**\n",
        "\n",
        "```json\n",
        "{\n",
        "  \"granter_classification\": \"United States or European Union\",\n",
        "  \"invention_type\": \"med_tech/computer_vision/crypto/other\",\n",
        "  \"entities\": {\n",
        "    \"issuer\": \"string\",\n",
        "    \"language\": \"string\",\n",
        "    \"publication_date\": \"YYYY-MM-DD\",\n",
        "    \"class_international\": [\"string\", \"string\", ...],\n",
        "    \"class_us\": [\"string\", \"string\", ...],\n",
        "    \"application_number\": \"string\",\n",
        "    \"filing_date\": \"YYYY-MM-DD\",\n",
        "    \"priority_date_eu\": \"YYYY-MM-DD\",\n",
        "    \"representative_eu\": [\"string\",\"string\",...],\n",
        "    \"applicant\": [\"string\", \"string\", ...],\n",
        "    \"inventor\": [\"string\", \"string\", ...],\n",
        "    \"title\": \"string\",\n",
        "    \"patent_number\": \"string\"\n",
        "   },\n",
        "  \"image_detections\": [\n",
        "        {\n",
        "          \"label\": \"figure_1\",\n",
        "          \"bounding_box\": [x1, y1, x2, y2]\n",
        "        },\n",
        "        {\n",
        "          \"label\": \"figure_2\",\n",
        "          \"bounding_box\": [x1, y1, x2, y2]\n",
        "        }\n",
        "     ]\n",
        "}\n",
        "```\n",
        "\n",
        "**Detailed Instructions:**\n",
        "\n",
        "1.  **Granter Classification:** Determine if the patent was granted by the \"United States\" or \"European Union\".  Populate the `granter_classification` field accordingly.\n",
        "\n",
        "2.  **Invention Type Classification:** Classify the type of invention described in the patent as one of the following and populate the `invention_type` field with the corresponding code:\n",
        "    *   Medical Tech:  `med_tech`\n",
        "    *   Computer Vision: `computer_vision`\n",
        "    *   Cryptography: `crypto`\n",
        "    *   Other: `other`\n",
        "\n",
        "3.  **Entity Extraction:**\n",
        "    *   Extract the following entities and populate the corresponding fields within the `entities` object.\n",
        "    *   If multiple entities exist for a field (e.g., multiple inventors, EU representatives, multiple US and international classes), represent them as a JSON array of strings.\n",
        "    *   Use the `YYYY-MM-DD` format for dates.\n",
        "    *   If a field is not present in the document, leave the value as null or an empty array as appropriate\n",
        "\n",
        "4.  **Image Object Detection:**\n",
        "    *   Detect the bounding boxes of all image figures in the patent document.\n",
        "    *   Represent each bounding box as a JSON array with the format `[x1, y1, x2, y2]` where:\n",
        "        *   `x1` and `y1` are the coordinates of the top-left corner of the bounding box.\n",
        "        *   `x2` and `y2` are the coordinates of the bottom-right corner of the bounding box.\n",
        "    *   Include a label (e.g., \"figure_1\", \"figure_2\", etc.) for each detected bounding box.\n",
        "    *   Populate the `image_detections` array with each detected bounding box. If no figures are detected, the image_detections array should be empty.\n",
        "\n",
        "**Input:**\n",
        "\n",
        "Provide the complete patent document text as input to this prompt.\n",
        "\n",
        "**Output:**\n",
        "\n",
        "The output should be a single JSON object following the specified format, containing all the extracted information.\n",
        "\n",
        "**Example (Illustrative, Not Complete):**\n",
        "\n",
        "```json\n",
        "{\n",
        "    \"granter_classification\": \"United States\",\n",
        "    \"invention_type\": \"med_tech\",\n",
        "    \"entities\": {\n",
        "        \"issuer\": \"United States Patent and Trademark Office\",\n",
        "        \"language\": \"en\",\n",
        "        \"publication_date\": \"2023-10-26\",\n",
        "        \"class_international\": [\"A61K 31/4709\", \"A61K 9/00\"],\n",
        "         \"class_us\": [\"514/250\"],\n",
        "        \"application_number\": \"17/744,843\",\n",
        "        \"filing_date\": \"2022-05-24\",\n",
        "        \"priority_date_eu\": null,\n",
        "        \"representative_eu\": [],\n",
        "        \"applicant\": [\"ABC Pharmacy Inc\"],\n",
        "        \"inventor\": [\"John Smith\", \"Jane Doe\"],\n",
        "        \"title\": \"Novel Formulation\",\n",
        "        \"patent_number\": \"US12345678\"\n",
        "    },\n",
        "    \"image_detections\": [\n",
        "      {\n",
        "        \"label\": \"figure_1\",\n",
        "        \"bounding_box\": [100, 50, 300, 250]\n",
        "      },\n",
        "       {\n",
        "         \"label\": \"figure_2\",\n",
        "         \"bounding_box\": [400, 100, 600, 300]\n",
        "       }\n",
        "    ]\n",
        "}\n",
        "```\n",
        "\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01048a7ef243"
      },
      "source": [
        "### Create BigQuery Dataset and Table for Batch Prediction\n",
        "\n",
        "NOTE: The location of the BigQuery dataset must be the same as the location for Vertex AI. (e.g. `us-central1`, not `us`)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d-KVqc67ajwf"
      },
      "outputs": [],
      "source": [
        "DATASET_NAME = \"patents_data_batch\"\n",
        "TABLE_NAME = \"patents_table\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e6de4b4be476"
      },
      "outputs": [],
      "source": [
        "!bq --location={LOCATION} mk --dataset \"{PROJECT_ID}:{DATASET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c403e052d071"
      },
      "source": [
        "Load `GenerateContentRequest` JSON into BigQuery table.\n",
        "\n",
        "This request takes a Google Cloud Storage path to a PDF, processes it with the prompt above, and returns a structured dictionary using [Controlled Generation](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/control-generated-output)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bbba7e8861af"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "\n",
        "def create_request_json(row) -> str:\n",
        "    return json.dumps(\n",
        "        {\n",
        "            \"contents\": [\n",
        "                {\n",
        "                    \"role\": \"user\",\n",
        "                    \"parts\": [\n",
        "                        {\"text\": PATENTS_PROMPT},\n",
        "                        {\n",
        "                            \"fileData\": {\n",
        "                                \"fileUri\": row[\"gcs_path\"],\n",
        "                                \"mimeType\": \"application/pdf\",\n",
        "                            }\n",
        "                        },\n",
        "                    ],\n",
        "                }\n",
        "            ],\n",
        "            \"systemInstruction\": {\n",
        "                \"parts\": [{\"text\": \"You are an expert at analyzing patent documents.\"}]\n",
        "            },\n",
        "            \"generationConfig\": {\n",
        "                \"temperature\": 0,\n",
        "                \"responseMimeType\": \"application/json\",\n",
        "                \"responseSchema\": {\n",
        "                    \"type\": \"OBJECT\",\n",
        "                    \"properties\": {\n",
        "                        \"granter_classification\": {\n",
        "                            \"type\": \"STRING\",\n",
        "                            \"enum\": [\"United States\", \"European Union\"],\n",
        "                        },\n",
        "                        \"invention_type\": {\n",
        "                            \"type\": \"STRING\",\n",
        "                            \"enum\": [\"med_tech\", \"computer_vision\", \"crypto\", \"other\"],\n",
        "                        },\n",
        "                        \"entities\": {\n",
        "                            \"type\": \"OBJECT\",\n",
        "                            \"properties\": {\n",
        "                                \"issuer\": {\"type\": \"STRING\"},\n",
        "                                \"language\": {\"type\": \"STRING\"},\n",
        "                                \"publication_date\": {\n",
        "                                    \"type\": \"STRING\",\n",
        "                                    \"format\": \"date\",\n",
        "                                },\n",
        "                                \"class_international\": {\n",
        "                                    \"type\": \"ARRAY\",\n",
        "                                    \"items\": {\"type\": \"STRING\"},\n",
        "                                },\n",
        "                                \"class_us\": {\n",
        "                                    \"type\": \"ARRAY\",\n",
        "                                    \"items\": {\"type\": \"STRING\"},\n",
        "                                },\n",
        "                                \"application_number\": {\"type\": \"STRING\"},\n",
        "                                \"filing_date\": {\"type\": \"STRING\", \"format\": \"date\"},\n",
        "                                \"priority_date_eu\": {\n",
        "                                    \"type\": \"STRING\",\n",
        "                                    \"format\": \"date\",\n",
        "                                    \"nullable\": True,\n",
        "                                },\n",
        "                                \"representative_eu\": {\n",
        "                                    \"type\": \"ARRAY\",\n",
        "                                    \"items\": {\"type\": \"STRING\"},\n",
        "                                },\n",
        "                                \"applicant\": {\n",
        "                                    \"type\": \"ARRAY\",\n",
        "                                    \"items\": {\"type\": \"STRING\"},\n",
        "                                },\n",
        "                                \"inventor\": {\n",
        "                                    \"type\": \"ARRAY\",\n",
        "                                    \"items\": {\"type\": \"STRING\"},\n",
        "                                },\n",
        "                                \"title\": {\"type\": \"STRING\"},\n",
        "                                \"patent_number\": {\"type\": \"STRING\"},\n",
        "                            },\n",
        "                            \"required\": [\n",
        "                                \"issuer\",\n",
        "                                \"language\",\n",
        "                                \"publication_date\",\n",
        "                                \"class_international\",\n",
        "                                \"class_us\",\n",
        "                                \"application_number\",\n",
        "                                \"filing_date\",\n",
        "                                \"applicant\",\n",
        "                                \"inventor\",\n",
        "                                \"title\",\n",
        "                                \"patent_number\",\n",
        "                            ],\n",
        "                        },\n",
        "                        \"image_detections\": {\n",
        "                            \"type\": \"ARRAY\",\n",
        "                            \"items\": {\n",
        "                                \"type\": \"OBJECT\",\n",
        "                                \"properties\": {\n",
        "                                    \"label\": {\"type\": \"STRING\"},\n",
        "                                    \"bounding_box\": {\n",
        "                                        \"type\": \"ARRAY\",\n",
        "                                        \"items\": {\"type\": \"NUMBER\"},\n",
        "                                        \"minItems\": 4,\n",
        "                                        \"maxItems\": 4,\n",
        "                                    },\n",
        "                                },\n",
        "                                \"required\": [\"label\", \"bounding_box\"],\n",
        "                            },\n",
        "                        },\n",
        "                    },\n",
        "                    \"required\": [\n",
        "                        \"granter_classification\",\n",
        "                        \"invention_type\",\n",
        "                        \"entities\",\n",
        "                        \"image_detections\",\n",
        "                    ],\n",
        "                },\n",
        "            },\n",
        "        }\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "29afb3b3fdcb"
      },
      "outputs": [],
      "source": [
        "df[\"request\"] = df.apply(create_request_json, axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "271aeda2bdaa"
      },
      "outputs": [],
      "source": [
        "pandas_gbq.to_gbq(df, f\"{DATASET_NAME}.{TABLE_NAME}\", project_id=PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XKpdIpbw6Ycv"
      },
      "source": [
        "### Batch Process Patent Documents with Gemini\n",
        "\n",
        "Batch Processing in Vertex AI will take in the BigQuery table with the requests and return the output in a new BigQuery table in the same dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c069acc57089"
      },
      "outputs": [],
      "source": [
        "batch_job = client.batches.create(\n",
        "    model=MODEL_ID,\n",
        "    src=f\"bq://{PROJECT_ID}.{DATASET_NAME}.{TABLE_NAME}\",\n",
        ")\n",
        "batch_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e67c90f1117e"
      },
      "source": [
        "Run the following block to get the latest status of the batch job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cf230deeb013"
      },
      "outputs": [],
      "source": [
        "batch_job = client.batches.get(name=batch_job.name)\n",
        "batch_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bf44c2b0bd92"
      },
      "source": [
        "Run the following block to continuously poll the status of the batch job until it completes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "defb5eb0b1c1"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "# Refresh the job until complete\n",
        "while batch_job.state == \"JOB_STATE_RUNNING\":\n",
        "    batch_job = client.batches.get(name=batch_job.name)\n",
        "    # print(batch_job)\n",
        "    time.sleep(5)\n",
        "\n",
        "# Check if the job succeeds\n",
        "if batch_job.state == \"JOB_STATE_SUCCEEDED\":\n",
        "    print(\"Job succeeded!\")\n",
        "else:\n",
        "    print(f\"Job failed: {batch_job.error}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acd872565dc0"
      },
      "source": [
        "Load the results from the destination BigQuery table."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ef9a49490eb"
      },
      "outputs": [],
      "source": [
        "results_df = pandas_gbq.read_gbq(\n",
        "    batch_job.dest.bigquery_uri.replace(\"bq://\", \"\"), project_id=PROJECT_ID\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ed26cf628fc3"
      },
      "source": [
        "Extract the fields from the controlled generation response."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ac2c0a3d98c"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "\n",
        "def flatten_response(response) -> dict | None:\n",
        "    try:\n",
        "        parsed_json = json.loads(\n",
        "            response[\"candidates\"][0][\"content\"][\"parts\"][0][\"text\"]\n",
        "        )\n",
        "\n",
        "        entities = parsed_json.get(\"entities\", {})\n",
        "        additional_fields = {\n",
        "            \"granter_classification\": parsed_json.get(\"granter_classification\"),\n",
        "            \"image_detections\": parsed_json.get(\"image_detections\"),\n",
        "            \"invention_type\": parsed_json.get(\"invention_type\"),\n",
        "        }\n",
        "\n",
        "        return {**entities, **additional_fields}\n",
        "    except (KeyError, IndexError, json.JSONDecodeError) as e:\n",
        "        print(f\"Error processing response: {e}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "results_df = results_df.join(\n",
        "    pd.json_normalize(results_df[\"response\"].apply(flatten_response))\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "da7bda40a26e"
      },
      "outputs": [],
      "source": [
        "results_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UkIB1pek6_Y7"
      },
      "source": [
        "### Compare Results with Ground Truth\n",
        "\n",
        "We can pull the ground truth data from BigQuery into a DataFrame to compare the results. Note that the data is in a slightly different format, so a direct comparison is not straightforward."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzDxgh74mKVy"
      },
      "outputs": [],
      "source": [
        "query = \"\"\"\n",
        "SELECT\n",
        "  COALESCE(t1.gcs_path, t2.gcs_path, t3.gcs_path) AS gcs_path,\n",
        "  t1.issuer,\n",
        "  t1.language,\n",
        "  t1.publication_date,\n",
        "  t1.class_international,\n",
        "  t1.class_us,\n",
        "  t1.application_number,\n",
        "  t1.filing_date,\n",
        "  t1.priority_date_eu,\n",
        "  t1.representative_line_1_eu,\n",
        "  t1.applicant_line_1,\n",
        "  t1.inventor_line_1,\n",
        "  t1.title_line_1,\n",
        "  t1.number,\n",
        "  t2.x_relative_min,\n",
        "  t2.y_relative_min,\n",
        "  t2.x_relative_max,\n",
        "  t2.y_relative_max,\n",
        "  t3.invention_type\n",
        "FROM\n",
        "  `bigquery-public-data.labeled_patents.extracted_data` AS t1\n",
        "FULL OUTER JOIN\n",
        "  `bigquery-public-data.labeled_patents.figures` AS t2\n",
        "ON\n",
        "  t1.gcs_path = t2.gcs_path\n",
        "FULL OUTER JOIN\n",
        "  `bigquery-public-data.labeled_patents.invention_types` AS t3\n",
        "ON\n",
        "  COALESCE(t1.gcs_path, t2.gcs_path) = t3.gcs_path;\n",
        "\"\"\"\n",
        "\n",
        "query_job = bq_client.query(query)\n",
        "ground_truth_df = bq_client.query(query).result().to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "This notebook demonstrates the capabilities of using Gemini to extract structured information from complex documents, simplifying document understanding pipelines without needing to train custom models.\n",
        "\n",
        "**Key Improvements:**\n",
        "\n",
        "- **Comprehensive Extraction**: Unlike AutoML, which often limited field extraction to the first line (e.g., applicant, inventor, class US), Gemini accurately extracts full text and all listed values for these fields.\n",
        "- **Simplified Workflow**: AutoML required four separate models and four requests per document to complete the tasks. Gemini consolidates this into a single request.\n",
        "- **No Custom Model Training**: Gemini is a pre-trained model, eliminating the need for use case-specific training, saving both time and resources."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "patents_understanding.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
